##
rm(list=ls())
gc()


library(data.table)
library(ggplot2)
library(scales)
library(dplyr)
library(stringr)
source('theme_jake.R')

### Figure 2 -- PARTISAN TRANSITION MATRICES

# Here, we output the transition matrices figures detailing the number and percentage,
# of voters of each year1 party that are of each party in year2, for each linked sample.



# Load in analysis data
data = rbindlist(lapply(c('panel-analysis-2008-2012.csv.gz',
                          'panel-analysis-2012-2016.csv.gz',
                          'panel-analysis-2016-2020.csv.gz'), FUN = function(x){
                            d=fread(x, select = c('Party_year1','DemDiff','RepDiff'))[,Years:=str_replace(str_replace(x, '.csv.gz', ''),'panel-analysis-','')]
                            
                            
                            return(d)
                          }))
gc()

# Set party codes for year1 and year2
data[!Party_year1%in%c('Democrat','Republican'),Party_year1:='Non-Partisan']
data[DemDiff==1|(Party_year1=='Democrat' & DemDiff==0),Party_year2:='Democrat']
data[RepDiff==1|(Party_year1=='Republican' & RepDiff==0),Party_year2:='Republican']

data[!Party_year2%in%c('Democrat','Republican'),Party_year2:='Non-Partisan']
data[,DemDiff:=NULL]
data[,RepDiff:=NULL]

# Make year1 and year2 variables
data[,Year1:=substr(Years,1,4)]
data[,Year2:=substr(Years,6,9)]

# Create 2016-2020 matrix
d = data[Years=='2016-2020']
p = d[,list(Total = .N),by=c('Party_year1','Party_year2')]
p[,N_1:=sum(Total),by = 'Party_year1']

p[,prob:= round(Total/N_1,3)]


p[,color := case_when(Party_year1 == 'Republican' & Party_year2 == 'Republican' ~ '1',
                      Party_year1 == 'Democrat' & Party_year2 == 'Democrat' ~ '2',
                      Party_year1 == 'Non-Partisan' & Party_year2 == 'Non-Partisan' ~ '3',
                      Party_year1 == 'Republican' & Party_year2 == 'Non-Partisan'|Party_year1=='Non-Partisan'&Party_year2=='Republican' ~ '4',
                      Party_year1 == 'Democrat' & Party_year2 == 'Non-Partisan'|Party_year1=='Non-Partisan'&Party_year2=='Democrat' ~ '5',
                      Party_year1 == 'Republican' & Party_year2 == 'Democrat'|Party_year1=='Democrat'&Party_year2=='Republican' ~ '6'
                      
)]  
p[Party_year1!=Party_year2 & Party_year1 %in% c('Democrat', 'Republican'), a:='c']
p[Party_year1!=Party_year2 & (Party_year2 =='Non-Partisan' | Party_year1 =='Non-Partisan'), a:='b']
p[Party_year1==Party_year2, a:='a']
p[,Party_year1:=paste0('2016\n',Party_year1)]
p[,Party_year2:=paste0('2020\n',Party_year2)]


g1=ggplot(p, aes(y = factor(Party_year1), x = factor(Party_year2), fill = color, label =prob))+
  geom_tile(aes(alpha=a))+
  scale_fill_manual(values=c('red','blue', 'grey','pink','lightblue','purple'))+
  scale_alpha_manual(values = c(.5,.2,.1))+
  geom_text( size = 8, color = 'black') +
  geom_text(aes(label = paste0('(n = ',comma(Total),')')),nudge_y=-.15, size = 5, color = 'black') +
  #scale_x_discrete(position = "top") +
  coord_equal() +
  geom_vline(xintercept = 1:3 - 0.5, colour = "white", size = 1.5) +
  geom_hline(yintercept = 1:3 - 0.5, colour = "white", size = 1.5) +
  ylab("") +  
  coord_equal() +
  xlab(paste(''))+
  ylab(paste(''))+
  theme_jake_tile2()+
  theme(axis.text.y=element_text(angle=90,hjust=.5), text = element_text(size=24))

# store output
p1620=p



## Make 2012-2016 matrix

d = data[Years=='2012-2016']
gc()


p = d[,list(Total = .N),by=c('Party_year1','Party_year2')]
p[,N_1:=sum(Total),by = 'Party_year1']

p[,prob:= round(Total/N_1,3)]

p[,color := ifelse(Party_year1 == 'Democrat', '1',ifelse(Party_year2=='Republican', '3', '2'))]

p[Party_year1!=Party_year2 & Party_year1 %in% c('Democrat', 'Republican'), a:='c']
p[Party_year1!=Party_year2 & (Party_year2 =='Non-Partisan' | Party_year1 =='Non-Partisan'), a:='b']
p[Party_year1==Party_year2, a:='a']

p[,color := case_when(Party_year1 == 'Republican' & Party_year2 == 'Republican' ~ '1',
                      Party_year1 == 'Democrat' & Party_year2 == 'Democrat' ~ '2',
                      Party_year1 == 'Non-Partisan' & Party_year2 == 'Non-Partisan' ~ '3',
                      Party_year1 == 'Republican' & Party_year2 == 'Non-Partisan'|Party_year1=='Non-Partisan'&Party_year2=='Republican' ~ '4',
                      Party_year1 == 'Democrat' & Party_year2 == 'Non-Partisan'|Party_year1=='Non-Partisan'&Party_year2=='Democrat' ~ '5',
                      Party_year1 == 'Republican' & Party_year2 == 'Democrat'|Party_year1=='Democrat'&Party_year2=='Republican' ~ '6'
                      
)]  
p[Party_year1!=Party_year2 & Party_year1 %in% c('Democrat', 'Republican'), a:='c']
p[Party_year1!=Party_year2 & (Party_year2 =='Non-Partisan' | Party_year1 =='Non-Partisan'), a:='b']
p[Party_year1==Party_year2, a:='a']
p[,Party_year1:=paste0('2012\n',Party_year1)]
p[,Party_year2:=paste0('2016\n',Party_year2)]


g2=ggplot(p, aes(y = factor(Party_year1), x = factor(Party_year2), fill = color, label =prob))+
  geom_tile(aes(alpha=a))+
  scale_fill_manual(values=c('red','blue', 'grey','pink','lightblue','purple'))+
  scale_alpha_manual(values = c(.5,.2,.1))+
  geom_text( size = 8, color = 'black') +
  geom_text(aes(label = paste0('(n = ',comma(Total),')')),nudge_y=-.15, size = 5, color = 'black') +
  #scale_x_discrete(position = "top") +
  coord_equal() +
  geom_vline(xintercept = 1:3 - 0.5, colour = "white", size = 1.5) +
  geom_hline(yintercept = 1:3 - 0.5, colour = "white", size = 1.5) +
  ylab("") +  
  coord_equal() +
  xlab(paste(''))+
  ylab(paste(''))+
  theme_jake_tile2()+
  theme(axis.text.y=element_text(angle=90,hjust=.5), text = element_text(size=24))

# store output
p1216=p


## Make 2008 - 2012 matrix
d = data[Years=='2008-2012']

p = d[,list(Total = .N),by=c('Party_year1','Party_year2')]
p[,N_1:=sum(Total),by = 'Party_year1']

p[,prob:= round(Total/N_1,3)]

p[,color := ifelse(Party_year1 == 'Democrat', '1',ifelse(Party_year2=='Republican', '3', '2'))]

p[Party_year1!=Party_year2 & Party_year1 %in% c('Democrat', 'Republican'), a:='c']
p[Party_year1!=Party_year2 & (Party_year2 =='Non-Partisan' | Party_year1 =='Non-Partisan'), a:='b']
p[Party_year1==Party_year2, a:='a']

p[,color := case_when(Party_year1 == 'Republican' & Party_year2 == 'Republican' ~ '1',
                      Party_year1 == 'Democrat' & Party_year2 == 'Democrat' ~ '2',
                      Party_year1 == 'Non-Partisan' & Party_year2 == 'Non-Partisan' ~ '3',
                      Party_year1 == 'Republican' & Party_year2 == 'Non-Partisan'|Party_year1=='Non-Partisan'&Party_year2=='Republican' ~ '4',
                      Party_year1 == 'Democrat' & Party_year2 == 'Non-Partisan'|Party_year1=='Non-Partisan'&Party_year2=='Democrat' ~ '5',
                      Party_year1 == 'Republican' & Party_year2 == 'Democrat'|Party_year1=='Democrat'&Party_year2=='Republican' ~ '6'
                      
)]  
p[Party_year1!=Party_year2 & Party_year1 %in% c('Democrat', 'Republican'), a:='c']
p[Party_year1!=Party_year2 & (Party_year2 =='Non-Partisan' | Party_year1 =='Non-Partisan'), a:='b']
p[Party_year1==Party_year2, a:='a']

p[,Party_year1:=paste0('2008\n',Party_year1)]
p[,Party_year2:=paste0('2012\n',Party_year2)]


g3=ggplot(p, aes(y = factor(Party_year1), x = factor(Party_year2), fill = color, label =prob))+
  geom_tile(aes(alpha=a))+
  scale_fill_manual(values=c('red','blue', 'grey','pink','lightblue','purple'))+
  scale_alpha_manual(values = c(.5,.2,.1))+
  geom_text( size = 8, color = 'black') +
  geom_text(aes(label = paste0('(n = ',comma(Total),')')),nudge_y=-.15, size = 5, color = 'black') +
  #scale_x_discrete(position = "top") +
  coord_equal() +
  geom_vline(xintercept = 1:3 - 0.5, colour = "white", size = 1.5) +
  geom_hline(yintercept = 1:3 - 0.5, colour = "white", size = 1.5) +
  ylab("") +  
  coord_equal() +
  xlab(paste(''))+
  ylab(paste(''))+
  theme_jake_tile2()+
  theme(axis.text.y=element_text(angle=90,hjust=.5), text = element_text(size=24))
library(cowplot)

ggsave(plot = plot_grid(g3,g2,g1,nrow=1, labels = c('2008-2012', '2012-2016', '2016-2020'), hjust=-3.1), file = paste0('figures/Fig2.png'),width=18,height=6)

